#!/usr/bin/env python
# -*- coding: utf-8 -*-
from datetime import datetime

import numpy as np
from scipy.stats import boxcox
import pandas as pd
from django.db.models import Q
from matplotlib import pyplot as plt
from prophet import Prophet
from sklearn.preprocessing import StandardScaler

from web.constants.datetime_format import DatetimeFormat
from web.manager.log_manager import LogManager
from web.models import CommodityFutureDateData
from web.models.stock_index_week import StockIndexWeek
from web.task.base_task import BaseTask
from web.util.datetime_util import DatetimeUtil
from web.util.picture_util import PictureUtil

Logger = LogManager.get_logger(__name__)


class ProphetTask(BaseTask):
    """
    Prophet 是由 Facebook 开发的一个开源时间序列预测库，设计考虑了业务场景中的时间序列特点，如季节性变化、假日效应和趋势变化。
    Prophet 特别适合处理日级别（或以上频率）的时间序列数据，并且在处理缺失数据和异常值方面表现出色。
    参考资料：https://blog.csdn.net/m0_64336780/article/details/138803481
    """

    def do_task(self, code: str, begin_date: str, end_date: str):
        """
        运行程序
        """

        # 准备数据
        begin_date: datetime.date = DatetimeUtil.str_to_datetime(begin_date, DatetimeFormat.Date_Format)
        end_date: datetime.date = DatetimeUtil.str_to_datetime(end_date, DatetimeFormat.Date_Format)
        stock_index_week_queryset = StockIndexWeek.objects.filter(
            Q(end_date__range=[begin_date, end_date]) & Q(code_=code)).order_by('end_date')
        # stock_index_week_queryset = CommodityFutureDateData.objects.filter(
        #     Q(transaction_date__range=[begin_date, end_date]) & Q(code=code)).order_by('transaction_date')

        if stock_index_week_queryset is not None and len(stock_index_week_queryset) > 0:
            # queryset转换为dataframe
            stock_index_week_dataframe = pd.DataFrame(list(stock_index_week_queryset.values()))

            # 创建直方图
            close_price_list = list(stock_index_week_dataframe['close_price'])
            PictureUtil.create_histogram_picture('收盘价', 'prophet_close_price.png', close_price_list, False)

            # 取end_date和close_price两列
            stock_index_week_dataframe = stock_index_week_dataframe[['transaction_date', 'close_price']]

            # 列重命名
            stock_index_week_dataframe.rename(columns={'transaction_date': 'ds', 'close_price': 'y'}, inplace=True)

            # 转换格式
            stock_index_week_dataframe['ds'] = stock_index_week_dataframe['ds'].apply(
                lambda x: str(x))
            stock_index_week_dataframe['y'] = stock_index_week_dataframe['y'].apply(pd.to_numeric)

            stock_index_week_dataframe['y'] = np.log(stock_index_week_dataframe['y'])

            # 创建直方图
            close_price_list = list(stock_index_week_dataframe['y'])
            PictureUtil.create_histogram_picture('收盘价', 'prophet_close_price_2.png', close_price_list, False)

            # Box-Cox转换
            stock_index_week_ndarray = stock_index_week_dataframe.values
            # temp = stock_index_week_ndarray[:, 1:].astype('float')
            # stock_index_week_ndarray[:, 1], _ = boxcox(temp.flatten() + 1, lmbda=None, alpha=None)

            # 归一化
            # stock_index_week_dataframe = pd.DataFrame(data=stock_index_week_ndarray[0:, 0:], columns=['ds', 'y'])
            # scaler = StandardScaler()
            # stock_index_week_dataframe[['y']] = scaler.fit_transform(stock_index_week_dataframe[['y']])

            # 创建直方图
            close_price_list = list(stock_index_week_dataframe['y'])
            PictureUtil.create_histogram_picture('收盘价', 'prophet_close_price_3.png', close_price_list, False)

            # 开始预测
            m = Prophet()
            m.fit(stock_index_week_dataframe)
            future = m.make_future_dataframe(periods=100, freq='D')
            # future.tail()
            forecast = m.predict(future)
            forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail()
            m.plot(forecast)
            m.plot_components(forecast)
            plt.show()
            pass

    def do_task2(self, code: str, begin_date: str, end_date: str):
        # 用 csv 建立数据框，转为对数坐标
        # df = pd.read_csv('C:/mywork/temp/example_wp_log_peyton_manning.csv')
        # df['y'] = np.log(df['y'])
        # df.head()

        # 准备数据
        begin_date: datetime.date = DatetimeUtil.str_to_datetime(begin_date, DatetimeFormat.Date_Format)
        end_date: datetime.date = DatetimeUtil.str_to_datetime(end_date, DatetimeFormat.Date_Format)
        stock_index_week_queryset = StockIndexWeek.objects.filter(
            Q(end_date__range=[begin_date, end_date]) & Q(code_=code)).order_by('end_date')

        df = pd.DataFrame(list(stock_index_week_queryset.values()))
        df = df[['end_date', 'close_price']]
        df['ds'] = df['end_date']
        df['y'] = np.log(df['close_price'].apply(pd.to_numeric))
        df.head()

        # 创建直方图
        close_price_list = list(df['y'])
        PictureUtil.create_histogram_picture('收盘价', 'example_wp_log_peyton_manning_1.png', close_price_list, False)

        # 将数据框导入 Prophet 模型
        m = Prophet()
        m.fit(df)
        # 预测，将结果放入 forecast 数据框
        future = m.make_future_dataframe(periods=20, freq='W')
        forecast = m.predict(future)
        forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail()
        # 按 forecast 数据框绘图
        m.plot(forecast)
        m.plot_components(forecast)
        x1 = forecast['ds']
        y1 = forecast['yhat']
        y2 = forecast['yhat_lower']
        y3 = forecast['yhat_upper']
        plt.plot(x1, y1)
        plt.plot(x1, y2)
        plt.plot(x1, y3)
        plt.show()
        pass


if __name__ == '__main__':
    prophet_task = ProphetTask()
    # prophet_task.do_task("V", "20221001", "20240101")
    prophet_task.do_task2("000001", "20220101", "20240801")
